#include <cusp/gallery/poisson.h>
#include <cusp/io/matrix_market.h>
#include <cusp/multiply.h>

#include <cusp/coo_matrix.h>
#include <cusp/csr_matrix.h>
#include <cusp/permutation_matrix.h>

#include <iostream>
#include <stdio.h>

#include "../timer.h"

template <typename MatrixType1, typename MatrixType2>
float time_row_permutation(const MatrixType1& A,
                           const MatrixType2& B)
{
    unsigned int N = 10;

    timer t;

    for(unsigned int i = 0; i < N; i++)
    {
        MatrixType2 C;
        cusp::multiply(A, B, C);
    }

    return t.milliseconds_elapsed() / N;
}

template <typename MatrixType1, typename MatrixType2>
float time_col_permutation(const MatrixType1& A,
                           const MatrixType2& B)
{
    unsigned int N = 10;

    timer t;

    for(unsigned int i = 0; i < N; i++)
    {
        MatrixType1 C;
        cusp::multiply(A, B, C);
    }

    return t.milliseconds_elapsed() / N;
}

int main(int argc, char ** argv)
{
    cudaSetDevice(0);

    typedef int    IndexType;
    typedef float  ValueType;

    cusp::csr_matrix<IndexType, ValueType, cusp::host_memory> A;

    if (argc == 1)
    {
        // no input file was specified, generate an example
        cusp::gallery::poisson5pt(A, 500, 500);
    }
    else if (argc == 2)
    {
        // an input file was specified, read it from disk
        cusp::io::read_matrix_market_file(A, argv[1]);
    }

    cusp::permutation_matrix<IndexType, cusp::host_memory> B(A.num_rows);
    thrust::sequence(B.values.rbegin(), B.values.rend());

    std::cout << "Input matrix A has shape (" << A.num_rows << "," << A.num_cols << ") and " << A.num_entries << " entries" << "\n";
    std::cout << "             B has shape (" << B.num_rows << "," << B.num_cols << ") and " << B.num_entries << " entries" << "\n\n";
   
    printf("Host Sparse Permutation Time (milliseconds per permutation)\n");
    printf("    ROW PERM    | %9.2f\n"  , time_row_permutation(B,A));
    printf("    COL PERM    | %9.2f\n\n", time_col_permutation(A,B));

    cusp::coo_matrix<IndexType, ValueType, cusp::device_memory> C(A);
    cusp::permutation_matrix<IndexType, cusp::device_memory> D(B);
    
    printf("Device Sparse Permutation Time (milliseconds per permutation)\n");
    printf("    ROW PERM    | %9.2f\n"  , time_row_permutation(D,C));
    printf("    COL PERM    | %9.2f\n\n", time_col_permutation(C,D));
   
    return 0;
}

